import math, os, sys
from vtk import *



sample_atoms = (('C', (0,0,0)), ('C', (0,0,1)),
                ('H', (.5,0,1)), ('HO', (.5,.5,1)),
                ('O', (0,.5,.5)), ('Z',(1,1,1)))

carbon_sphere = {'radius' : .15, 'color' : (.2,.2,.2) }
silicon_sphere = {'radius' : .2, 'color' : (1,.2,.2) }
dopant_sphere = {'radius' : .1, 'color' : (.2, 1, .1) }
unknown_sphere = {'radius' : .1, 'color' : (1, 1, 1) }
spheresByType = { 'C' : carbon_sphere, 'H' : silicon_sphere, 'O' :
dopant_sphere,
            'HO' : silicon_sphere, 'unknown' : unknown_sphere }
            
sphereResolution = {'theta' : 20, 'phi' : 20 }


class AtomShow:
  def __init__(self):
    self.polyData = None
    self.typeToIdx = None
    self.CreatePipeline()
    self.ConfigureSpheres(spheresByType)
    
  def __CreateSphere(self, sphereProps):
    sphere = vtkSphereSource()
    sphere.SetRadius(sphereProps["radius"])
    sphere.SetThetaResolution(sphereResolution["theta"])
    sphere.SetPhiResolution(sphereResolution["phi"])
    return sphere.GetOutput()
  
  def GetActor(self):
    return self.glyphActor
  
  
  def CreateAtomPolyData(self, atomSet):
    createPts = vtkPoints()
    createPts.SetDataTypeToFloat()
    createIndices = vtkIdTypeArray()
    createCellArray = vtkCellArray()
    
    for atom in atomSet:
      ptIdx = createPts.InsertNextPoint(atom[1])
      if self.typeToIdx.has_key(atom[0]):
        typeIdx = self.typeToIdx[atom[0]]
      else:
        typeIdx = self.typeToIdx['unknown']
      createIndices.InsertNextValue(typeIdx)
      createCellArray.InsertNextCell(1)
      createCellArray.InsertCellPoint(ptIdx)
      
    newPolyData = vtkPolyData()
    newPolyData.SetPoints(createPts)
    newPolyData.GetPointData().SetScalars(createIndices)
    newPolyData.SetVerts(createCellArray)

    return newPolyData  
  
  def AddAtoms(self, atomPolyData):
    self.RegisterPolydataInPipeline(atomPolyData)
  
  def AddLookupTableColor(self, index, color):
    self.lut.SetTableValue(index, color[0], color[1], color[2], 1)

  def CreatePipeline(self):
    self.glyphs = vtkGlyph3D()
    self.lut = vtkLookupTable()
    self.glyphMapper = vtkPolyDataMapper()
    self.glyphActor = vtkActor()

  def RegisterPolydataInPipeline(self, newPolyData):
    self.glyphs.SetInput(newPolyData)
    if self.polyData and self.polyData is not newPolyData:
      del self.polyData
    self.polyData = newPolyData

  def ConfigureSpheres(self, sphereTypes):
    self.typeToIdx = {}
    sphereIdx = 0
    for sphereType in sphereTypes.keys():
      self.typeToIdx[sphereType] = sphereIdx
      sphereIdx = sphereIdx + 1
    sphereCnt = sphereIdx
    
    for sourceType in sphereTypes.keys():
      self.glyphs.SetSource(self.typeToIdx[sourceType],
                      self.__CreateSphere(sphereTypes[sourceType]))
    self.glyphs.SetScaleModeToDataScalingOff()
    self.glyphs.SetIndexModeToScalar()
    self.glyphs.SetRange(0,sphereCnt-1)
    
    self.lut.SetNumberOfTableValues(sphereCnt)
    for sourceType in sphereTypes.keys():
      idx = self.typeToIdx[sourceType]
      sphereColor = sphereTypes[sourceType]["color"]
      self.AddLookupTableColor(idx, sphereColor)
    
    self.glyphMapper.SetInput(self.glyphs.GetOutput())
    self.glyphMapper.SetLookupTable(self.lut)
    self.glyphMapper.SetScalarRange(0,sphereCnt-1)
    
    self.glyphActor.SetMapper(self.glyphMapper)
    self.glyphActor.GetProperty().SetOpacity(1.0)
    self.glyphActor.GetProperty().SetInterpolationToGouraud()
    
    
    
if __name__ == "__main__":
  actorCreator = AtomShow()
  actorCreator.AddAtoms(sample_atoms)
  actorCreator.AddAtoms(sample_atoms)
  actorCreator.GetActor()
